using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace JufGame
{
	public class ActivateFunc : BaseLayer
	{
        private bool isLoadLastOutput;
        public ActivateFunc(ComputeShader shader, int batch, int inputCount, bool isLoadLastOutput) : base(shader, batch, inputCount)
        {
            this.isLoadLastOutput = isLoadLastOutput;
            outputBuffer = new ComputeBuffer(batch * inputCount, sizeof(float));
        }

        public override void Forward(ComputeBuffer input)
        {
			compute.SetInt("batch", batch);
			compute.SetInt("inputCount", inputCount);
            base.Forward(input);
        }

        /// <summary>
        /// 基于compute shader的反向传播
        /// </summary>
        /// <param name="dout">求导数据</param>
        public override void Backward(ComputeBuffer dOutputBuffer)
		{
			compute.SetInt("batch", batch);
			compute.SetInt("inputCount", inputCount);
            if(isLoadLastOutput)
            {
                compute.SetBuffer(kernelBackward, "outputBuffer", outputBuffer);
            }
			base.Backward(dOutputBuffer);
		}
	}
}